"""
AkShare数据源工具模块
"""
import logging
import pandas as pd
from datetime import datetime, timedelta
import akshare as ak

logger = logging.getLogger(__name__)


class AkShareClient:
    """
    AkShare客户端封装，用于从AkShare获取股票数据
    """
    
    @staticmethod
    def get_stock_list():
        """
        获取股票列表信息
        
        Returns:
            pandas.DataFrame: 包含股票代码、名称等基本信息的DataFrame
        """
        try:
            # 使用akshare获取A股所有股票列表
            stock_info_df = ak.stock_info_a_code_name()
            # 转换列名以匹配模型字段
            stock_info_df.columns = ['symbol', 'name']
            # 添加ts_code字段 (格式: 000001.SZ)
            stock_info_df['ts_code'] = stock_info_df['symbol'].apply(
                lambda x: x + '.SH' if x.startswith(('6', '9')) else x + '.SZ'
            )
            return stock_info_df
        except Exception as e:
            logger.error(f"获取股票列表失败: {str(e)}")
            return pd.DataFrame()
    
    @staticmethod
    def get_stock_industry():
        """
        获取股票行业分类信息
        
        Returns:
            pandas.DataFrame: 包含股票代码、所属行业的DataFrame
        """
        try:
            # 使用akshare获取股票行业分类
            industry_df = ak.stock_industry_sina()
            # 转换列名以匹配模型字段
            return industry_df
        except Exception as e:
            logger.error(f"获取股票行业分类失败: {str(e)}")
            return pd.DataFrame()
    
    @staticmethod
    def get_daily_data(symbol, start_date=None, end_date=None):
        """
        获取股票日线行情数据
        
        Args:
            symbol (str): 股票代码，如 '000001'（不含市场后缀）
            start_date (str, optional): 开始日期，格式 'YYYYMMDD'，默认为近30天
            end_date (str, optional): 结束日期，格式 'YYYYMMDD'，默认为今天
            
        Returns:
            pandas.DataFrame: 包含日线行情数据的DataFrame
        """
        try:
            # 如果未提供日期，默认获取最近30天数据
            if not end_date:
                end_date = datetime.now().strftime('%Y%m%d')
            if not start_date:
                start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
                
            # 使用akshare获取股票日线数据
            daily_df = ak.stock_zh_a_hist(symbol=symbol, period="daily", 
                                          start_date=start_date, end_date=end_date, 
                                          adjust="qfq")
            
            # 转换列名以匹配模型字段
            daily_df = daily_df.rename(columns={
                '日期': 'trade_date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'vol',
                '成交额': 'amount',
                '振幅': 'amplitude', 
                '涨跌幅': 'pct_chg',
                '涨跌额': 'change',
                '换手率': 'turnover'
            })
            
            # 转换日期格式
            daily_df['trade_date'] = pd.to_datetime(daily_df['trade_date'])
            
            return daily_df
        except Exception as e:
            logger.error(f"获取股票 {symbol} 日线数据失败: {str(e)}")
            return pd.DataFrame()
    
    @staticmethod
    def get_financial_indicator(ts_code, period=20231231):
        """
        获取财务指标数据
        
        Args:
            ts_code (str): 股票代码，格式 '000001.SZ'
            period (int): 财报期，如 20231231 表示2023年年报
            
        Returns:
            pandas.DataFrame: 包含财务指标数据的DataFrame
        """
        try:
            # 提取股票代码前6位
            symbol = ts_code.split('.')[0]
            
            # 使用akshare获取财务指标
            # 注意：akshare的接口可能需要根据实际情况调整
            financial_df = ak.stock_financial_report_sina(stock=symbol, symbol="资产负债表")
            
            # 其他报表也可以获取
            profit_df = ak.stock_financial_report_sina(stock=symbol, symbol="利润表")
            cash_flow_df = ak.stock_financial_report_sina(stock=symbol, symbol="现金流量表")
            
            # TODO: 处理财务数据，提取需要的指标
            
            return financial_df
        except Exception as e:
            logger.error(f"获取股票 {ts_code} 财务指标失败: {str(e)}")
            return pd.DataFrame()
    
    @staticmethod
    def get_index_data(symbol="000001", period="daily", start_date=None, end_date=None):
        """
        获取指数数据
        
        Args:
            symbol (str): 指数代码，默认上证指数 '000001'
            period (str): 周期，'daily'表示日线
            start_date (str, optional): 开始日期，格式 'YYYYMMDD'
            end_date (str, optional): 结束日期，格式 'YYYYMMDD'
            
        Returns:
            pandas.DataFrame: 包含指数数据的DataFrame
        """
        try:
            # 如果未提供日期，默认获取最近30天数据
            if not end_date:
                end_date = datetime.now().strftime('%Y%m%d')
            if not start_date:
                start_date = (datetime.now() - timedelta(days=30)).strftime('%Y%m%d')
                
            # 使用akshare获取指数数据
            if symbol.startswith('0'):  # 上证指数
                index_df = ak.stock_zh_index_daily(symbol=symbol)
            elif symbol.startswith('3'):  # 深证指数
                index_df = ak.stock_zh_index_daily(symbol=symbol)
            else:
                raise ValueError(f"不支持的指数代码: {symbol}")
            
            # 过滤日期范围
            index_df = index_df[(index_df.index >= start_date) & (index_df.index <= end_date)]
            
            return index_df
        except Exception as e:
            logger.error(f"获取指数 {symbol} 数据失败: {str(e)}")
            return pd.DataFrame() 